from __future__ import annotations

import inspect
import logging
from typing import Any

from ConfigSpace import (
    CategoricalHyperparameter,
    Constant,
    OrdinalHyperparameter,
    UniformFloatHyperparameter,
    UniformIntegerHyperparameter,
)

from pruna.algorithms.base.pruna_base import PrunaAlgorithmBase
from pruna.algorithms.base.registry import AlgorithmRegistry
from pruna.algorithms.base.tags import AlgorithmTag
from pruna.config.hyperparameters import UnconstrainedHyperparameter

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")


def generate_rst_anchor(algorithm_name: str) -> str:
    """
    Generate the RST anchor that will be created from a header.

    RST generates anchors by converting headers to lowercase and replacing
    non-alphanumeric characters with hyphens.

    Parameters
    ----------
    algorithm_name: str
        The algorithm name to convert to an anchor.

    Returns
    -------
    str
        The RST anchor that will be generated from the header.
    """
    # RST wraps algorithm names in backticks in the header: ``algorithm_name``
    # The anchor generation converts this to lowercase and replaces non-alphanumeric chars
    header_text = f"``{algorithm_name}``"
    # Convert to lowercase and replace non-alphanumeric with hyphens
    anchor = header_text.lower().replace("`", "").replace("_", "-")
    return anchor


def generate_algorithm_desc(obj: PrunaAlgorithmBase, name_suffix: str = "") -> str:
    """
    Generate a Markdown description of a Pruna algorithm from its instance.

    Parameters
    ----------
    obj: PrunaAlgorithmBase
        The instance of the Pruna algorithm to generate the description for.
    name_suffix: str, optional
        A suffix to add to the algorithm name in the header.

    Returns
    -------
    str
        A Markdown description of the Pruna algorithm.
    """
    header = get_header_and_description(obj, name_suffix)
    references_str = get_references(obj)
    required_inputs_str = get_required_inputs(obj)
    compatible_devices_str = get_compatible_devices(obj)
    compatible_algorithms_str = get_compatible_algorithms(obj)
    required_install_str = get_required_install(obj)

    static_info = "\n".join(
        [
            f"| **References**: {references_str}.",
            f"| **Can be applied on**: {compatible_devices_str}.",
            f"| **Required**: {required_inputs_str}.",
            f"| **Compatible with**: {compatible_algorithms_str}.",
            f"| **Required install**: {required_install_str}."
            if required_install_str
            else "",
        ]
    )

    table_rows, hyperparameter_counter = get_table_rows(obj)

    # Only add the table if there are hyperparameters to show.
    if hyperparameter_counter > 0:
        table = format_grid_table(table_rows)
        markdown = f"{header}\n\n{static_info}\n\n{table}"
    else:
        markdown = f"{header}\n\n{static_info}"
    return markdown


def get_algorithm_description(cls: Any) -> str:
    """Extract the description segment from the docstring of a class."""
    doc = inspect.getdoc(cls)
    assert doc is not None, "No docstring found for class"
    lines = doc.splitlines()
    assert len(lines) > 2, "No description found in docstring"
    segment_lines = []
    for line in lines[2:]:
        if not line.strip():
            break
        segment_lines.append(line.strip())
    return " ".join(segment_lines)


def format_grid_table(rows: list[list[str]]) -> str:
    """Given a list of rows (each row is a list of cell strings), returns a grid table."""
    num_cols = len(rows[0])
    col_widths = [max(len(row[i]) for row in rows) for i in range(num_cols)]
    total_widths = [w + 2 for w in col_widths]

    horizontal_border = "+" + "+".join("-" * width for width in total_widths) + "+"
    header_line = (
        "|"
        + "|".join(" " + rows[0][i].ljust(col_widths[i]) + " " for i in range(num_cols))
        + "|"
    )
    header_separator = "+" + "+".join("=" * width for width in total_widths) + "+"

    data_lines = []
    for row in rows[1:]:
        row_line = (
            "|"
            + "|".join(" " + row[i].ljust(col_widths[i]) + " " for i in range(num_cols))
            + "|"
        )
        data_lines.append(row_line)
        data_lines.append(horizontal_border)

    return "\n".join([horizontal_border, header_line, header_separator] + data_lines)


def get_header_and_description(obj: PrunaAlgorithmBase, name_suffix: str = "") -> str:
    """Get the header and description of a Pruna algorithm."""
    algo = obj.algorithm_name
    description = get_algorithm_description(obj)
    header = f"``{algo}`` {name_suffix}\n{'~' * (len(algo) + len(name_suffix) + 5)}\n"
    header += f"{description}\n"
    return header


def get_references(obj: PrunaAlgorithmBase) -> str:
    """Get the references of a Pruna algorithm."""
    references = obj.references
    if references:
        return ", ".join(f"`{key} <{value}>`__" for key, value in references.items())
    return "None"


def get_required_inputs(obj: PrunaAlgorithmBase) -> str:
    """Get the required inputs of a Pruna algorithm."""
    required_inputs = []
    if obj.tokenizer_required:
        required_inputs.append("Tokenizer")
    if obj.processor_required:
        required_inputs.append("Processor")
    if obj.dataset_required:
        required_inputs.append("Dataset")
    return ", ".join(required_inputs) if required_inputs else "None"


def get_compatible_devices(obj: PrunaAlgorithmBase) -> str:
    """Get the compatible devices of a Pruna algorithm."""
    compatible_devices = []
    name_map = {
        "cpu": "CPU",
        "cuda": "CUDA",
        "mps": "MPS",
        "accelerate": "Accelerate distributed",
    }
    for device in obj.runs_on:
        compatible_devices.append(name_map.get(device, device))
    return ", ".join(compatible_devices) if compatible_devices else "None"


def get_compatible_algorithms(obj: PrunaAlgorithmBase) -> str:
    """Get the compatible algorithms of a Pruna algorithm."""
    compatible_algorithms = obj.get_compatible_algorithms()
    # Sort alphabetically by algorithm name only
    compatible_algorithms.sort(key=str.lower)
    # Format with intra-page RST links using correct anchors
    linked_algorithms = [
        f"`{a} <compression.html#{generate_rst_anchor(a)}>`__"
        for a in compatible_algorithms
    ]
    return ", ".join(linked_algorithms) if linked_algorithms else "None"


def get_required_install(obj: PrunaAlgorithmBase) -> str | None:
    """Get the required install of a Pruna algorithm."""
    required_install = obj.required_install
    base_install = "``pip install pruna[full]``"

    # Break the long command into multiple code blocks with line continuation for better RST rendering
    pro_install = (
        "``pip install pruna_pro[full]`` "
        "``--extra-index-url https://prunaai.pythonanywhere.com/`` "
        "``--extra-index-url https://pytorch-extension.intel.com/release-whl/stable/cpu/cn/``"
    )
    if required_install:
        if obj.algorithm_name == "gptq":
            return required_install
        if "pro" in required_install:
            required_install += f" or {pro_install}"
        else:
            required_install += f" or {base_install}"
        return required_install
    return None


def get_table_rows(obj: PrunaAlgorithmBase) -> tuple[list[list[str]], int]:
    """Get the table rows of a Pruna algorithm hyperparameter section."""
    rows = [["**Parameter**", "**Default**", "**Options**", "**Description**"]]
    hyperparameter_counter = 0
    for hp in obj.get_hyperparameters():
        if isinstance(hp, Constant) and not isinstance(hp, UnconstrainedHyperparameter):
            continue
        param_name = f"``{obj.algorithm_name}_{hp.name}``"
        # Fix: Check if meta is None before accessing it
        description = hp.meta.get("desc", "") if hp.meta is not None else ""
        if isinstance(hp, (UniformFloatHyperparameter, UniformIntegerHyperparameter)):
            default = str(hp.default_value)
            values = f"Range {hp.lower} to {hp.upper}"
        elif isinstance(hp, OrdinalHyperparameter):
            choices = ", ".join(str(v) for v in hp.sequence[:-1])
            values = f"{choices} or {hp.sequence[-1]}"
            default = str(hp.default_value)
        elif isinstance(hp, CategoricalHyperparameter):
            values = ", ".join(str(v) for v in hp.choices)
            default = str(hp.default_value)
        elif isinstance(hp, UnconstrainedHyperparameter):
            default = str(hp.value)
            values = "Unconstrained"
        else:
            raise ValueError(f"Unsupported hyperparameter type: {type(hp)}")
        rows.append([param_name, default, values, description])
        hyperparameter_counter += 1
    return rows, hyperparameter_counter


def generate_compatibility_table() -> str:
    """Generate a reStructuredText list-table showing algorithm compatibility."""
    all_algorithms = AlgorithmRegistry.get_all_algorithms()

    all_algorithms.sort(key=lambda algo: algo.algorithm_name.lower())

    lines = [
        ".. list-table:: Algorithm Compatibility Matrix",
        "   :header-rows: 1",
        "",
        "   * - Algorithm",
        "     - Compatible With",
    ]

    for algo in all_algorithms:
        compatibles = []
        for lst in algo.get_compatible_algorithms():
            compatibles.extend(lst)
        compatibles.sort(key=str.lower)
        # Fix: Use correct RST anchor generation
        compat_text = (
            ", ".join(
                f"`{c} <compression.html#{generate_rst_anchor(c)}>`__"
                for c in compatibles
            )
            if compatibles
            else "—"
        )
        lines.append(f"   * - ``{algo.algorithm_name}``\n     - {compat_text}")

    return "\n".join(lines)


if __name__ == "__main__":
    # Collect all algorithms into a single file.
    with open("compression.rst", "w") as f:
        # First write the compatibility matrix
        f.write(generate_compatibility_table())
        f.write("\n\n")

        # Write algorithm descriptions
        for tag in AlgorithmTag.__members__.values():
            for algorithm_name in AlgorithmRegistry.get_algorithms_by_tag(tag):
                algorithm = AlgorithmRegistry[algorithm_name]
                f.write(generate_algorithm_desc(algorithm))
                f.write("\n\n")

    logger.info(
        "Generated compression.rst with compatibility matrix and algorithm descriptions"
    )
    logger.info("Docs successfully written to docs/user_manual/")
